﻿module TypeRules
open TypeEnvironment
open MaybeMonad
open Ast

let rec private check_literal (literal : Ast.Literal) : TypeCheckState = 
  maybe{
    match literal with
    | Int _   -> return Type.Integer
    | Str _   -> return Type.String
    | Bool _  -> return Type.Boolean
    | Null    -> return (Type.Identifier "$null$")
  }

and check_expr (expr : Ast.Expr) : TypeCheckState = 
  maybe{
    match expr with
    | Expr.Binary (e1, op, e2) -> 
      let! t1 = check_expr e1
      let! t2 = check_expr e2
      match op with
      | Binary.And | Binary.Or -> 
        if t1 = Type.Boolean && t2 = Type.Boolean then 
          return Type.Boolean
        else
          return! fail (sprintf "Type mismatch; found %A, %A; espected %A %A" t1 t2 Type.Boolean Type.Boolean)
      | Binary.Sum sum when t1 = Type.String && t2 = Type.String ->
        do sum := false
        return Type.String
      | Binary.Div | Binary.Minus | Binary.Mod | Binary.Mul | Binary.Sum _ -> 
        if t1 = Type.Integer && t2 = Type.Integer then 
          return Type.Integer
        else
          return! fail (sprintf "Type mismatch; found %A, %A; espected %A %A" t1 t2 Type.Integer Type.Integer)
      | Binary.Greater | Binary.GreaterEq | Binary.Min | Binary.MinEq -> 
        if t1 = Type.Integer && t2 = Type.Integer then 
          return Type.Boolean
        else
          return! fail (sprintf "Type mismatch; found %A, %A; espected %A %A" t1 t2 Type.Integer Type.Integer)
      | Binary.Eq | Binary.NotEq -> 
        if t1 = t2 then 
          return Type.Boolean
        else
          match t1,t2 with
          | ty1, ty2 when ty1 = Type.Identifier "$null$" ->
            match ty2 with
            | Type.String | Type.Array _ | Type.Identifier _ ->
              return Type.Boolean
            | _ -> 
              return! fail (sprintf "Type mismatch; type %A not match %A" t1 t2)
          | ty1, ty2 when ty2 = Type.Identifier "$null$" ->
            match ty1 with
            | Type.String | Type.Array _ | Type.Identifier _ ->
              return Type.Boolean
            | _ -> 
              return! fail (sprintf "Type mismatch; type %A not match %A" t1 t2)
//          | Type.Identifier n1, Type.Identifier n2 ->
//            if n1 = "$null$" || n2 = "$null$" then
//              return Type.Boolean
//            else
//              return! fail (sprintf "Type mismatch; type %A not match %A" t1 t2)
//          | Type.String, Type.Identifier id when id = "$null$" ->
//            return Type.Boolean
//          | Type.Identifier id, Type.String when id = "$null$" ->
//            return Type.Boolean
          | _ ->
            return! fail (sprintf "Type mismatch; type %A not match %A" t1 t2)
    | Expr.Unary (op, e1) ->
      let! t1 = check_expr e1
      match op with
      | Unary.Negative -> 
        if t1 = Type.Integer then
          return Type.Integer
        else
          return! fail (sprintf "Type mismatch; found %A; espected %A" t1 Type.Integer)
      | Unary.Not -> 
        if t1 = Type.Boolean then
          return Type.Boolean
        else
          return! fail (sprintf "Type mismatch; found %A; espected %A" t1 Type.Boolean)
    | Expr.Literal l ->
      return! check_literal l
    | Expr.This -> 
      let! container_class = get_container_class_name
      return Type.Identifier container_class
    | Expr.New name -> 
      do! reduce_class ((=) name) (sprintf "type %s doesn't exist" name)
      return Type.Identifier name
    | Expr.NewArray (ty, e1) -> 
      do! check_type ty
      let! t1 = check_expr e1
      if t1 = Type.Integer then
        return Type.Array ty
      else
        return! fail (sprintf "Type mismatch; found %A; espected %A" t1 Type.Integer)
    | Expr.Length expr -> 
      let! t1 = check_expr expr
      match t1 with
      | Type.Array _ ->
        return Type.Integer
      | _ ->
        return! fail (sprintf "Type mismatch; found %A; espected Array" t1)
    | Expr.Call call -> 
      let! ret = check_call call
      match ret with
      | ReturnType.Type ty -> 
        return ty
      | ReturnType.Void ->
        return! fail "expression call must be non void"
    | Expr.Location loc -> 
      return! check_loc loc
  }

and check_call (call : Ast.Call) =
  maybe{
    let (Call.Call (meth, exprs)) = call
    let! (ret, args) = check_method meth
    let rec check_arg_tys (args : Type list) (exprs : Expr list) = 
      maybe{
        match args, exprs with
        | [], [] -> return ()
        | _, [] -> return! fail "missing arguments"
        | [], _ -> return! fail "too much arguments"
        | (arg::args),(expr::exprs) -> 
          let! texpr = check_expr expr            
          do! type_compare arg texpr
      }
    do! check_arg_tys args exprs
    return ret
  }

and check_method (meth : Ast.Method) = 
  maybe{
    match meth with
    | Method.Method name ->
      return! try_find_method name
    | Method.Select ((Expr.Location (Location.Id loc)), name, ty) when loc = "Library" -> 
      do ty := None
      match name with
      | "println" ->
        return (ReturnType.Void, [Type.String])
      | "print" ->
        return (ReturnType.Void, [Type.String])
      | "printi" ->
        return (ReturnType.Void, [Type.Integer])
      | "printb" ->
        return (ReturnType.Void, [Type.Boolean])
      | "readi" ->
        return (ReturnType.Type Type.Integer, [])
      | "readln" ->
        return (ReturnType.Type Type.String, [])
      | "eof" ->
        return (ReturnType.Type Type.Boolean, [])
      | "stoi" ->
        return (ReturnType.Type Type.Integer, [Type.String; Type.Integer])
      | "itos" ->
        return (ReturnType.Type Type.String, [Type.Integer])
      | "stoa" ->
        return (ReturnType.Type (Type.Array Type.Integer), [Type.String])
      | "aots" ->
        return (ReturnType.Type Type.String, [Type.Array Type.Integer])
      | "random" ->
        return (ReturnType.Type Type.Integer, [Type.Integer])
      | "time" ->
        return (ReturnType.Type Type.Integer, [])
      | "exit" ->
        return (ReturnType.Type Type.Integer, [Type.Integer])
      | name ->
        return! fail (sprintf "Method %s not in library" name)      
    | Method.Select (expr, name, ty) ->
      let! t1 = check_expr expr
      do ty := Some t1
      match t1 with
      | Type.Identifier ty ->
        return! try_find_method_with ty name        
      | _ -> 
        return! fail (name + "doesn't exist")        
  }

and check_loc (loc : Ast.Location) : TypeCheckState =
  maybe{
    match loc with
    | Location.Id name ->
      let! ty = find_var name

      return ty
    | Location.Member (expr, name, ty) ->
      let! texpr = check_expr expr
      do ty := Some texpr
      match texpr with
      | Type.Identifier ty ->
        return! try_find_var_with ty name        
      | _ -> 
        return! fail (name + "doesn't exist")
    | Location.Nth (e1, e2) ->
      let! t1 = check_expr e1
      let! t2 = check_expr e2
      match t1,t2 with
      | Type.Array ty, Type.Integer ->
        return ty
      | ty, Type.Integer ->
        return! fail (sprintf "type %A is not Array" ty)
      | Type.Array _, ty ->
        return! fail (sprintf "type %A is not Integer" ty)
      | _,_ -> 
        return! fail (sprintf "type %A is not Array and type %A is not Integer" t1 t2)
  }

and check_type (ty : Ast.Type) : Maybe<Unit> = 
  maybe{
    match ty with
    | Type.Identifier name ->
      do! reduce_class ((=) name) (sprintf "type %s doesn't exist" name)
    | Type.Array ty ->
      do! check_type ty
    | _ ->
      return ()
  }

and check_stmt (stmt : Ast.Stmt) = 
  maybe{
    match stmt with
    | Stmt.Assign (loc, expr, line) -> 
      do! set_line line
      let! is_arg = is_arg loc
      if is_arg |> not then
        let! t_loc = check_loc loc
        let! t_expr = check_expr expr
        do! type_compare t_loc t_expr
      else
        return! fail "arguments are not assignable"
    | Stmt.Call (call, line) -> 
      do! set_line line
      let! _ = check_call call
      return()
    | Stmt.Break line
    | Stmt.Continue line -> 
      do! set_line line
      do! check_is_while
    | Stmt.Return (Some expr, line) -> 
      do! set_line line
      let! t_expr = check_expr expr
      let! return_ty = get_metod_return_type
      match return_ty with
      | ReturnType.Type ty ->
        do! type_compare ty t_expr
      | _ ->
        return! fail (sprintf "found %A instead of void" t_expr)
    | Stmt.Return (None, line) -> 
      do! set_line line
      let! return_ty = get_metod_return_type
      if return_ty <> ReturnType.Void then
        return! fail (sprintf "expected %A instead of void" return_ty)
      else
        return ()
    | Stmt.If (cond, body, None, line) ->
      do! set_line line
      let! ty_cond = check_expr cond
      if ty_cond = Type.Boolean then
        return! check_stmt body
      else
        return! fail (sprintf "expected bool; found %A" ty_cond)
    | Stmt.If (cond, body, Some else_body, line) ->
      do! set_line line
      let! ty_cond = check_expr cond
      if ty_cond = Type.Boolean then
        do! check_stmt body
        return! check_stmt else_body
      else
        return! fail (sprintf "expected bool; found %A" ty_cond)
    | Stmt.While (cond, body, line) ->
      do! set_line line
      let! ty_cond = check_expr cond
      if ty_cond = Type.Boolean then
        do! enter_while
        do! check_stmt body
        do! exit_while
      else
        return! fail (sprintf "expected bool; found %A" ty_cond)
    | Stmt.Block block ->
      do! check_block block      
  }

and check_var_decl (var : Ast.VarDecl) = 
  maybe{
    let check_assoc (ty:Type) (name:id, expr:Expr option) =
      maybe{
        do! check_field_name name
        match expr with
        | None ->
          return ()
        | Some expr ->
          let! ty_expr = check_expr expr
          do! type_compare ty ty_expr
      }
    match var with
    | VarDecl (ty,(id, expr), others, line) ->
      do! set_line line
      do! check_type ty
      do! check_assoc ty (id, expr)
      do! add_field_symbol id ty
      for (id, expr) in others do
        do! check_assoc ty (id, expr)
        do! add_field_symbol id ty
  }

and check_block (block : Ast.Block) = 
  maybe{
    let (Block.Block(vars, stmts, line)) = block
    do! set_line line
    do! add_block_env block
    for var in vars do
      do! check_var_decl var
    for stmt in stmts do
      do! check_stmt stmt
    do! remove_env
  }

and check_arg (arg : Ast.Arg) = 
  maybe{
    let (Arg(ty, name)) = arg
    do! check_type ty
    do! check_field_name name
    do! add_field_symbol name ty
  }

and check_return_type (ret : Ast.ReturnType) = 
  maybe{
    match ret with
    | Void ->
      return ()
    | ReturnType.Type t ->
      return! check_type t
  }

and check_field_or_method_decl (f_or_m : Ast.FieldOrMethodDecl) = 
  maybe{
    match f_or_m with
    | FieldOrMethodDecl.FieldDecl (ty, name, others, line) ->
      do! set_line line
      do! check_type ty
//      do! check_name name
//      for id in others do
//        do! check_name name
//        do! add_field_symbol name ty
    | FieldOrMethodDecl.MethodDecl (ret, name, args, block, line) ->
      do! set_line line
      do! check_return_type ret
//      do! check_name name
      do! reset_alias_counter
      do! add_method_env (ret, name, args, block)
      for arg in args do
        do! check_arg arg
      do! check_block block
      do! remove_env      
  }

and check_class_decl (c : Ast.ClassDecl) = 
  maybe{
    let (ClassDecl (name, inherited, fields)) = c
//    do! check_name name
    if inherited.IsSome then 
      do! reduce_class ((=) inherited.Value) (sprintf "class %s doesn't exist" inherited.Value)
    else
      do! maybe{return ()}
    
    do! add_class_env name

    for fieldOrMethod in fields do 
      do! check_field_or_method_decl fieldOrMethod
    do! remove_env
  }

and type_compare t1 t2 : Maybe<unit> = 
  maybe{
    match t1, t2 with
    | ty1, ty2 when ty1 = (Type.Identifier "$null$") ->
      match ty2 with
      | Type.String | Type.Array _ | Type.Identifier _ ->
        return ()
      | _ ->
        return! fail (sprintf "expected %A; found %A " t1 t2)
    | ty1, ty2 when ty2 = (Type.Identifier "$null$") ->
      match ty1 with
      | Type.String | Type.Array _ | Type.Identifier _ ->
        return ()
      | _ ->
        return! fail (sprintf "expected %A; found %A " t1 t2)
    | Type.Identifier name1, Type.Identifier name2 ->
      let! inheriteds = get_inherited_list_with name2
      if inheriteds |> Seq.exists (fun c -> c.Name = name1) then
        return ()
      else
        return! fail (sprintf "expected %A; found %A " name1 name2)
    | Type.Array t1, Type.Array t2 ->
      return! type_compare t1 t2
    | _ when t1 = t2 ->
      return ()
    | _ -> 
      return! fail (sprintf "expected %A; found %A " t1 t2)
  }

let check_program (prg : Ast.Program) = 
  maybe{
    for ``class decl ;-P`` in prg do
      do! check_class_decl ``class decl ;-P``
  }